"""
This is the main file for the runtime client.
It is responsible for executing actions received from OpenHands backend and producing observations.

NOTE: this will be executed inside the docker sandbox.
"""
from starlette.background import BackgroundTask
import argparse
import asyncio
import base64
import io
import json
import mimetypes
import os
import re
import shutil
import tempfile
import time
import traceback
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Iterator
from zipfile import ZipFile

from fastapi import Depends, FastAPI, HTTPException, Request, UploadFile
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.security import APIKeyHeader
from openhands_aci.utils.diff import get_diff
from pydantic import BaseModel
from starlette.exceptions import HTTPException as StarletteHTTPException
from uvicorn import run

from openhands.core.schema.action import ActionType
from openhands.core.logger import openhands_logger as logger
from openhands.events.action import (
    Action,
    BrowseInteractiveAction,
    BrowseURLAction,
    CmdRunAction,
    FileReadAction,
    FileWriteAction,
    FileViewAction,
    IPythonRunCellAction,
)
from openhands.events.event import FileEditSource, FileReadSource
from openhands.events.observation import (
    CmdOutputObservation,
    ErrorObservation,
    FileEditObservation,
    FileReadObservation,
    FileWriteObservation,
    FileViewObservation,
    IPythonRunCellObservation,
    Observation,
)
from openhands.events.serialization import event_from_dict, event_to_dict
from openhands.runtime.browser import browse
from openhands.runtime.browser.browser_env import BrowserEnv
from openhands.runtime.plugins import ALL_PLUGINS, JupyterPlugin, Plugin, VSCodePlugin
from openhands.runtime.utils.bash import BashSession
from openhands.runtime.utils.files import insert_lines, read_lines
from openhands.runtime.utils.runtime_init import init_user_and_working_directory
from openhands.runtime.utils.system import check_port_available
from openhands.runtime.utils.system_stats import get_system_stats
from openhands.utils.async_utils import call_sync_from_async, wait_all


class ActionRequest(BaseModel):
    action: dict


ROOT_GID = 0
INIT_COMMANDS = [
    'git config --global user.name "openhands" && git config --global user.email "openhands@all-hands.dev" && alias git="git --no-pager"',
]

SESSION_API_KEY = os.environ.get('SESSION_API_KEY')
api_key_header = APIKeyHeader(name='X-Session-API-Key', auto_error=False)


def verify_api_key(api_key: str = Depends(api_key_header)):
    if SESSION_API_KEY and api_key != SESSION_API_KEY:
        raise HTTPException(status_code=403, detail='Invalid API Key')
    return api_key


class ActionExecutor:
    """ActionExecutor is running inside docker sandbox.
    It is responsible for executing actions received from OpenHands backend and producing observations.
    """

    def __init__(
            self,
            plugins_to_load: list[Plugin],
            work_dir: str,
            username: str,
            user_id: int,
            browsergym_eval_env: str | None,
    ) -> None:
        self.plugins_to_load = plugins_to_load
        self._initial_pwd = work_dir
        self.username = username
        self.user_id = user_id
        _updated_user_id = init_user_and_working_directory(
            username=username, user_id=self.user_id, initial_pwd=work_dir
        )
        if _updated_user_id is not None:
            self.user_id = _updated_user_id

        self.bash_session = BashSession(
            work_dir=work_dir,
            username=username,
        )

        self.lock = asyncio.Lock()
        self.plugins: dict[str, Plugin] = {}
        self.browser = BrowserEnv(browsergym_eval_env)
        self.start_time = time.time()
        self.last_execution_time = self.start_time

    @property
    def initial_pwd(self):
        return self._initial_pwd

    async def ainit(self):
        await wait_all(
            (self._init_plugin(plugin) for plugin in self.plugins_to_load),
            timeout=30,
        )

        # This is a temporary workaround
        # TODO: refactor AgentSkills to be part of JupyterPlugin
        # AFTER ServerRuntime is deprecated
        if 'agent_skills' in self.plugins and 'jupyter' in self.plugins:
            obs = await self.run_ipython(
                IPythonRunCellAction(
                    code='from openhands.runtime.plugins.agent_skills.agentskills import *\n'
                )
            )
            logger.debug(f'AgentSkills initialized: {obs}')

        await self._init_bash_commands()
        logger.debug('Runtime client initialized.')

    async def _init_plugin(self, plugin: Plugin):
        await plugin.initialize(self.username)
        self.plugins[plugin.name] = plugin
        logger.debug(f'Initializing plugin: {plugin.name}')

        if isinstance(plugin, JupyterPlugin):
            await self.run_ipython(
                IPythonRunCellAction(
                    code=f'import os; os.chdir("{self.bash_session.pwd}")'
                )
            )

    async def _init_bash_commands(self):
        logger.debug(f'Initializing by running {len(INIT_COMMANDS)} bash commands...')
        for command in INIT_COMMANDS:
            action = CmdRunAction(command=command)
            action.timeout = 300
            logger.debug(f'Executing init command: {command}')
            obs = await self.run(action)
            assert isinstance(obs, CmdOutputObservation)
            logger.debug(
                f'Init command outputs (exit code: {obs.exit_code}): {obs.content}'
            )
            assert obs.exit_code == 0

        logger.debug('Bash init commands completed')

    async def run_action(self, action, stream_callback=None) -> Observation:
        async with self.lock:
            action_type = action.action
            logger.debug(f'Running action:\n{action}, function: {action_type}')
            if action_type == ActionType.RUN:
                observation = await getattr(self, action_type)(action, stream_callback)
            else:
                observation = await getattr(self, action_type)(action)
            logger.debug(f'Action output:\n{observation}')
            return observation

    async def stop_cmd_action(self):
        self.bash_session.intercept()

    async def run_stream(self, action: CmdRunAction):
        """以流式方式执行命令，并实时返回输出结果。"""
        async for obs in self.bash_session.run_stream(action):
            yield obs

    async def run(
            self, action: CmdRunAction, stream_callback=None
    ) -> CmdOutputObservation | ErrorObservation:
        action.stream_callback = stream_callback
        obs = await call_sync_from_async(self.bash_session.run, action)
        return obs

    async def run_ipython(self, action: IPythonRunCellAction) -> Observation:
        if 'jupyter' in self.plugins:
            _jupyter_plugin: JupyterPlugin = self.plugins['jupyter']  # type: ignore
            # This is used to make AgentSkills in Jupyter aware of the
            # current working directory in Bash
            jupyter_pwd = getattr(self, '_jupyter_pwd', None)
            if self.bash_session.pwd != jupyter_pwd:
                logger.debug(
                    f'{self.bash_session.pwd} != {jupyter_pwd} -> reset Jupyter PWD'
                )
                reset_jupyter_pwd_code = (
                    f'import os; os.chdir("{self.bash_session.pwd}")'
                )
                _aux_action = IPythonRunCellAction(code=reset_jupyter_pwd_code)
                _reset_obs: IPythonRunCellObservation = await _jupyter_plugin.run(
                    _aux_action
                )
                logger.debug(
                    f'Changed working directory in IPython to: {self.bash_session.pwd}. Output: {_reset_obs}'
                )
                self._jupyter_pwd = self.bash_session.pwd

            obs: IPythonRunCellObservation = await _jupyter_plugin.run(action)
            obs.content = obs.content.rstrip()
            matches = re.findall(
                r'<oh_aci_output_[0-9a-f]{32}>(.*?)</oh_aci_output_[0-9a-f]{32}>',
                obs.content,
                re.DOTALL,
            )
            if matches:
                results: list[str] = []
                if len(matches) == 1:
                    # Use specific actions/observations types
                    match = matches[0]
                    try:
                        result_dict = json.loads(match)
                        if result_dict.get('path'):  # Successful output
                            if (
                                    result_dict['new_content'] is not None
                            ):  # File edit commands
                                diff = get_diff(
                                    old_contents=result_dict['old_content']
                                                 or '',  # old_content is None when file is created
                                    new_contents=result_dict['new_content'],
                                    filepath=result_dict['path'],
                                )
                                return FileEditObservation(
                                    content=diff,
                                    path=result_dict['path'],
                                    old_content=result_dict['old_content'],
                                    new_content=result_dict['new_content'],
                                    prev_exist=result_dict['prev_exist'],
                                    impl_source=FileEditSource.OH_ACI,
                                    formatted_output_and_error=result_dict[
                                        'formatted_output_and_error'
                                    ],
                                )
                            else:  # File view commands
                                return FileReadObservation(
                                    content=result_dict['formatted_output_and_error'],
                                    path=result_dict['path'],
                                    impl_source=FileReadSource.OH_ACI,
                                )
                        else:  # Error output
                            results.append(result_dict['formatted_output_and_error'])
                    except json.JSONDecodeError:
                        # Handle JSON decoding errors if necessary
                        results.append(
                            f"Invalid JSON in 'openhands-aci' output: {match}"
                        )
                else:
                    for match in matches:
                        try:
                            result_dict = json.loads(match)
                            results.append(result_dict['formatted_output_and_error'])
                        except json.JSONDecodeError:
                            # Handle JSON decoding errors if necessary
                            results.append(
                                f"Invalid JSON in 'openhands-aci' output: {match}"
                            )

                # Combine the results (e.g., join them) or handle them as required
                obs.content = '\n'.join(str(result) for result in results)

            if action.include_extra:
                obs.content += (
                    f'\n[Jupyter current working directory: {self.bash_session.pwd}]'
                )
                obs.content += f'\n[Jupyter Python interpreter: {_jupyter_plugin.python_interpreter_path}]'
            return obs
        else:
            raise RuntimeError(
                'JupyterRequirement not found. Unable to run IPython action.'
            )

    def _resolve_path(self, path: str, working_dir: str) -> str:
        filepath = Path(path)
        if not filepath.is_absolute():
            return str(Path(working_dir) / filepath)
        return str(filepath)

    async def read(self, action: FileReadAction) -> Observation:
        if action.impl_source == FileReadSource.OH_ACI:
            return await self.run_ipython(
                IPythonRunCellAction(
                    code=action.translated_ipython_code,
                    include_extra=False,
                )
            )

        # NOTE: the client code is running inside the sandbox,
        # so there's no need to check permission
        working_dir = self.bash_session.workdir
        filepath = self._resolve_path(action.path, working_dir)
        try:
            if filepath.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
                with open(filepath, 'rb') as file:
                    image_data = file.read()
                    encoded_image = base64.b64encode(image_data).decode('utf-8')
                    mime_type, _ = mimetypes.guess_type(filepath)
                    if mime_type is None:
                        mime_type = 'image/png'  # default to PNG if mime type cannot be determined
                    encoded_image = f'data:{mime_type};base64,{encoded_image}'

                return FileReadObservation(path=filepath, content=encoded_image)
            elif filepath.lower().endswith('.pdf'):
                with open(filepath, 'rb') as file:
                    pdf_data = file.read()
                    encoded_pdf = base64.b64encode(pdf_data).decode('utf-8')
                    encoded_pdf = f'data:application/pdf;base64,{encoded_pdf}'
                return FileReadObservation(path=filepath, content=encoded_pdf)
            elif filepath.lower().endswith(('.mp4', '.webm', '.ogg')):
                with open(filepath, 'rb') as file:
                    video_data = file.read()
                    encoded_video = base64.b64encode(video_data).decode('utf-8')
                    mime_type, _ = mimetypes.guess_type(filepath)
                    if mime_type is None:
                        mime_type = 'video/mp4'  # default to MP4 if MIME type cannot be determined
                    encoded_video = f'data:{mime_type};base64,{encoded_video}'

                return FileReadObservation(path=filepath, content=encoded_video)

            with open(filepath, 'r', encoding='utf-8') as file:
                lines = read_lines(file.readlines(), action.start, action.end)
        except FileNotFoundError:
            return ErrorObservation(
                f'File not found: {filepath}. Your current working directory is {working_dir}.'
            )
        except UnicodeDecodeError:
            return ErrorObservation(f'File could not be decoded as utf-8: {filepath}.')
        except IsADirectoryError:
            return ErrorObservation(
                f'Path is a directory: {filepath}. You can only read files'
            )

        code_view = ''.join(lines)
        return FileReadObservation(path=filepath, content=code_view)

    async def write(self, action: FileWriteAction) -> Observation:
        working_dir = self.bash_session.workdir
        filepath = self._resolve_path(action.path, working_dir)

        insert = action.content.split('\n')
        try:
            if not os.path.exists(os.path.dirname(filepath)):
                os.makedirs(os.path.dirname(filepath))

            file_exists = os.path.exists(filepath)
            if file_exists:
                file_stat = os.stat(filepath)
            else:
                file_stat = None

            mode = 'w' if not file_exists else 'r+'
            try:
                with open(filepath, mode, encoding='utf-8') as file:
                    if mode != 'w':
                        all_lines = file.readlines()
                        new_file = insert_lines(
                            insert, all_lines, action.start, action.end
                        )
                    else:
                        new_file = [i + '\n' for i in insert]

                    file.seek(0)
                    file.writelines(new_file)
                    file.truncate()

                # Handle file permissions
                if file_exists:
                    assert file_stat is not None
                    # restore the original file permissions if the file already exists
                    os.chmod(filepath, file_stat.st_mode)
                    os.chown(filepath, file_stat.st_uid, file_stat.st_gid)
                else:
                    # set the new file permissions if the file is new
                    os.chmod(filepath, 0o664)
                    os.chown(filepath, self.user_id, self.user_id)

            except FileNotFoundError:
                return ErrorObservation(f'File not found: {filepath}')
            except IsADirectoryError:
                return ErrorObservation(
                    f'Path is a directory: {filepath}. You can only write to files'
                )
            except UnicodeDecodeError:
                return ErrorObservation(
                    f'File could not be decoded as utf-8: {filepath}'
                )
        except PermissionError:
            return ErrorObservation(f'Malformed paths not permitted: {filepath}')
        return FileWriteObservation(content='', path=filepath)

    async def browse(self, action: BrowseURLAction) -> Observation:
        return await browse(action, self.browser)

    async def browse_interactive(self, action: BrowseInteractiveAction) -> Observation:
        return await browse(action, self.browser)

    async def view(self, action: FileViewAction) -> Observation:
        # 根据 action.head 决定使用 head 还是 cat 命令
        command = f"cat -n {action.path}"
        if action.head > 0:
            command += f" | head -n {action.head}"
        try:
            # 异步执行命令
            process = await asyncio.create_subprocess_shell(
                command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            stdout, stderr = await process.communicate()

            # 检查命令是否成功
            if process.returncode != 0:
                error_message = stderr.decode()
                print(f"Error: {error_message}")
                return ErrorObservation(content=error_message)

            # 返回文件内容
            content = stdout.decode()
            return FileViewObservation(path=action.path, content=content, head=action.head)

        except Exception as e:
            # 捕获其他异常
            error_message = f"Unexpected error: {str(e)}"
            print(error_message)
            return ErrorObservation(content=error_message)

        except Exception as e:
            # 捕获其他异常
            error_message = f"Unexpected error: {str(e)}"
            print(error_message)
            return ErrorObservation(content=error_message)

    def close(self):
        self.bash_session.close()
        self.browser.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('port', type=int, help='Port to listen on')
    parser.add_argument('--working-dir', type=str, help='Working directory')
    parser.add_argument('--plugins', type=str, help='Plugins to initialize', nargs='+')
    parser.add_argument(
        '--username', type=str, help='User to run as', default='openhands'
    )
    parser.add_argument('--user-id', type=int, help='User ID to run as', default=1000)
    parser.add_argument(
        '--browsergym-eval-env',
        type=str,
        help='BrowserGym environment used for browser evaluation',
        default=None,
    )
    # example: python client.py 8000 --working-dir /workspace --plugins JupyterRequirement
    args = parser.parse_args()
    os.environ['VSCODE_PORT'] = str(int(args.port) + 1)
    assert check_port_available(int(os.environ['VSCODE_PORT']))

    plugins_to_load: list[Plugin] = []
    if args.plugins:
        for plugin in args.plugins:
            if plugin not in ALL_PLUGINS:
                raise ValueError(f'Plugin {plugin} not found')
            plugins_to_load.append(ALL_PLUGINS[plugin]())  # type: ignore

    client: ActionExecutor | None = None


    @asynccontextmanager
    async def lifespan(app: FastAPI):
        global client
        client = ActionExecutor(
            plugins_to_load,
            work_dir=args.working_dir,
            username=args.username,
            user_id=args.user_id,
            browsergym_eval_env=args.browsergym_eval_env,
        )
        await client.ainit()
        yield
        # Clean up & release the resources
        client.close()


    app = FastAPI(lifespan=lifespan)


    # TODO below 3 exception handlers were recommended by Sonnet.
    # Are these something we should keep?
    @app.exception_handler(Exception)
    async def global_exception_handler(request: Request, exc: Exception):
        logger.exception('Unhandled exception occurred:')
        return JSONResponse(
            status_code=500,
            content={'detail': 'An unexpected error occurred. Please try again later.'},
        )


    @app.exception_handler(StarletteHTTPException)
    async def http_exception_handler(request: Request, exc: StarletteHTTPException):
        logger.error(f'HTTP exception occurred: {exc.detail}')
        return JSONResponse(status_code=exc.status_code, content={'detail': exc.detail})


    @app.exception_handler(RequestValidationError)
    async def validation_exception_handler(
            request: Request, exc: RequestValidationError
    ):
        logger.error(f'Validation error occurred: {exc}')
        return JSONResponse(
            status_code=422,
            content={'detail': 'Invalid request parameters', 'errors': exc.errors()},
        )


    @app.middleware('http')
    async def authenticate_requests(request: Request, call_next):
        if request.url.path != '/alive' and request.url.path != '/server_info':
            try:
                verify_api_key(request.headers.get('X-Session-API-Key'))
            except HTTPException as e:
                return e
        response = await call_next(request)
        return response

    @app.middleware("http")
    async def track_client_connections(request: Request, call_next):
        # 为请求添加一个断开连接检测方法
        request.state.client_disconnected = False
        
        async def is_disconnected():
            return request.state.client_disconnected
        
        request.is_disconnected = is_disconnected
        
        try:
            response = await call_next(request)
            return response
        except Exception as e:
            # 如果发生异常，可能是客户端断开连接
            request.state.client_disconnected = True
            raise e

    @app.get('/server_info')
    async def get_server_info():
        assert client is not None
        current_time = time.time()
        uptime = current_time - client.start_time
        idle_time = current_time - client.last_execution_time

        response = {
            'uptime': uptime,
            'idle_time': idle_time,
            'resources': get_system_stats(),
        }
        logger.info('Server info endpoint response: %s', response)
        return response


    @app.post('/stop_cmd_action')
    async def stop_cmd_action():
        logger.info("============停止当前的CmdAction==========")
        assert client is not None
        try:
            client.last_execution_time = time.time()
            
            logger.info("正在停止当前执行的命令")
            await client.stop_cmd_action()
            logger.info("命令已成功停止")
            
            return JSONResponse(
                content={"status": "success", "message": "命令已成功停止"},
                status_code=200
            )
        except Exception as e:
            logger.error(
                f'Error processing command: {str(e)}', exc_info=True, stack_info=True
            )
            raise HTTPException(
                status_code=500,
                detail=traceback.format_exc(),
            )


    @app.post('/execute_action')
    async def execute_action(action_request: ActionRequest):
        logger.info("============执行action==========")
        assert client is not None
        try:
            action = event_from_dict(action_request.action)
            if not isinstance(action, Action):
                raise HTTPException(status_code=400, detail='Invalid action type')
            client.last_execution_time = time.time()
            observation = await client.run_action(action)
            return event_to_dict(observation)
        except Exception as e:
            logger.error(
                f'Error processing command: {str(e)}', exc_info=True, stack_info=True
            )
            raise HTTPException(
                status_code=500,
                detail=traceback.format_exc(),
            )


    def _stream_error(message: str) -> Iterator[str]:
        """返回错误消息的流。"""
        yield f"data: {message}\n\n"


    @app.post("/execute_action_stream")
    async def execute_action_stream(action_request: ActionRequest, request: Request) -> StreamingResponse:
        """执行动作并以流式方式返回结果。

        这个端点专门用于长时间运行的命令，可以实时返回命令的输出结果。
        """
        logger.info("开始============流式返回执行CmdRunAction==========")
        assert client is not None
        try:
            action = event_from_dict(action_request.action)
            if not isinstance(action, CmdRunAction):
                return StreamingResponse(
                    _stream_error("只支持 CmdRunAction 类型的动作进行流式执行"),
                    media_type="text/event-stream"
                )

            client.last_execution_time = time.time()

            # 创建一个任务ID，用于跟踪此次执行
            task_id = str(time.time())
            
            async def stream_generator():
                try:
                    # 设置客户端断开连接的检测
                    disconnect_event = asyncio.Event()
                    
                    async for obs in client.run_stream(action):
                        # 首先检查是否已经设置了断开连接标志
                        if disconnect_event.is_set():
                            # 客户端已断开连接，尝试停止命令执行
                            logger.warning(f"客户端断开连接，尝试停止命令执行: {action.command}")
                            await client.stop_cmd_action()
                            break
                        
                        # 生成数据前检查连接状态
                        try:
                            # 使用更安全的方式检查客户端连接状态
                            is_disconnected = False
                            try:
                                is_disconnected = await request.is_disconnected()
                            except Exception as conn_err:
                                # 如果检查连接状态时出错，假设客户端已断开
                                logger.warning(f"检查客户端连接状态时出错，假设已断开: {str(conn_err)}")
                                is_disconnected = True
                            
                            if is_disconnected:
                                disconnect_event.set()
                                logger.warning(f"检测到客户端断开连接: {action.command}")
                                await client.stop_cmd_action()
                                break
                            # 只有在客户端仍然连接的情况下才发送数据
                            print("==============发送中间数据到客户端===============")
                            yield f"data: {json.dumps(event_to_dict(obs))}\n\n"
                            
                        except Exception as yield_err:
                            # 如果在生成数据时出错，可能是客户端已断开
                            logger.warning(f"生成数据时出错，可能是客户端已断开: {str(yield_err)}")
                            disconnect_event.set()
                            await client.stop_cmd_action()
                            break
                            
                except Exception as e:
                    logger.error(f"流生成器出错: {str(e)}", exc_info=True)
                    # 尝试发送错误信息，但可能会失败如果客户端已断开
                    try:
                        yield f"data: {json.dumps(event_to_dict(ErrorObservation(f'流处理出错: {str(e)}', error_id='STREAM_ERROR')))}\n\n"
                    except Exception:
                        pass
                    finally:
                        # 确保命令被终止，防止资源泄漏
                        try:
                            await client.stop_cmd_action()
                        except Exception as e:
                            logger.error(f"停止命令时出错: {str(e)}")
                    
                logger.debug("结束============流式返回执行CmdRunAction========")

            # 不再使用 background 参数，因为它可能导致异常
            return StreamingResponse(
                stream_generator(),
                media_type="text/event-stream"
            )
        except Exception as e:
            logger.error(
                f'Error processing command: {str(e)}', exc_info=True, stack_info=True
            )
            return StreamingResponse(
                _stream_error(f"执行命令时出错: {str(e)}"),
                media_type="text/event-stream"
            )

    async def handle_client_disconnect(action: CmdRunAction):
        """处理客户端断开连接的情况"""
        logger.warning(f"客户端断开连接，清理资源: {action.command}")
        try:
            # 确保命令被终止
            await client.stop_cmd_action()
        except Exception as e:
            logger.error(f"断开连接处理时出错: {str(e)}")


    @app.post('/upload_file')
    async def upload_file(
            file: UploadFile, destination: str = '/', recursive: bool = False
    ):
        assert client is not None

        try:
            # Ensure the destination directory exists
            if not os.path.isabs(destination):
                raise HTTPException(
                    status_code=400, detail='Destination must be an absolute path'
                )

            full_dest_path = destination
            if not os.path.exists(full_dest_path):
                os.makedirs(full_dest_path, exist_ok=True)

            if recursive or file.filename.endswith('.zip'):
                # For recursive uploads, we expect a zip file
                if not file.filename.endswith('.zip'):
                    raise HTTPException(
                        status_code=400, detail='Recursive uploads must be zip files'
                    )

                zip_path = os.path.join(full_dest_path, file.filename)
                with open(zip_path, 'wb') as buffer:
                    shutil.copyfileobj(file.file, buffer)

                # Extract the zip file
                shutil.unpack_archive(zip_path, full_dest_path)
                os.remove(zip_path)  # Remove the zip file after extraction

                logger.debug(
                    f'Uploaded file {file.filename} and extracted to {destination}'
                )
            else:
                # For single file uploads
                file_path = os.path.join(full_dest_path, file.filename)
                with open(file_path, 'wb') as buffer:
                    shutil.copyfileobj(file.file, buffer)
                logger.debug(f'Uploaded file {file.filename} to {destination}')

            return JSONResponse(
                content={
                    'filename': file.filename,
                    'destination': destination,
                    'recursive': recursive,
                },
                status_code=200,
            )

        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))


    @app.get('/download_files')
    async def download_file(path: str):
        logger.debug('Downloading files')
        try:
            if not os.path.isabs(path):
                raise HTTPException(
                    status_code=400, detail='Path must be an absolute path'
                )

            if not os.path.exists(path):
                raise HTTPException(status_code=404, detail='File not found')

            with tempfile.TemporaryFile() as temp_zip:
                with ZipFile(temp_zip, 'w') as zipf:
                    for root, _, files in os.walk(path):
                        for file in files:
                            file_path = os.path.join(root, file)
                            zipf.write(
                                file_path, arcname=os.path.relpath(file_path, path)
                            )
                temp_zip.seek(0)  # Rewind the file to the beginning after writing
                content = temp_zip.read()
                # Good for small to medium-sized files. For very large files, streaming directly from the
                # file chunks may be more memory-efficient.
                zip_stream = io.BytesIO(content)
                return StreamingResponse(
                    content=zip_stream,
                    media_type='application/zip',
                    headers={'Content-Disposition': f'attachment; filename={path}.zip'},
                )

        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))


    @app.get('/alive')
    async def alive():
        return {'status': 'ok'}


    # ================================
    # VSCode-specific operations
    # ================================

    @app.get('/vscode/connection_token')
    async def get_vscode_connection_token():
        assert client is not None
        if 'vscode' in client.plugins:
            plugin: VSCodePlugin = client.plugins['vscode']  # type: ignore
            return {'token': plugin.vscode_connection_token}
        else:
            return {'token': None}


    # ================================
    # File-specific operations for UI
    # ================================

    @app.post('/list_files')
    async def list_files(request: Request):
        """List files in the specified path.

        This function retrieves a list of files from the agent's runtime file store,
        excluding certain system and hidden files/directories.

        To list files:
        ```sh
        curl http://localhost:3000/api/list-files
        ```

        Args:
            request (Request): The incoming request object.
            path (str, optional): The path to list files from. Defaults to '/'.

        Returns:
            list: A list of file names in the specified path.

        Raises:
            HTTPException: If there's an error listing the files.
        """
        assert client is not None

        # get request as dict
        request_dict = await request.json()
        path = request_dict.get('path', None)

        # Get the full path of the requested directory
        if path is None:
            full_path = client.initial_pwd
        elif os.path.isabs(path):
            full_path = path
        else:
            full_path = os.path.join(client.initial_pwd, path)

        if not os.path.exists(full_path):
            # if user just removed a folder, prevent server error 500 in UI
            return []

        try:
            # Check if the directory exists
            if not os.path.exists(full_path) or not os.path.isdir(full_path):
                return []

            entries = os.listdir(full_path)

            # Separate directories and files
            directories = []
            files = []
            for entry in entries:
                # Remove leading slash and any parent directory components
                entry_relative = entry.lstrip('/').split('/')[-1]

                # Construct the full path by joining the base path with the relative entry path
                full_entry_path = os.path.join(full_path, entry_relative)
                if os.path.exists(full_entry_path):
                    is_dir = os.path.isdir(full_entry_path)
                    if is_dir:
                        # add trailing slash to directories
                        # required by FE to differentiate directories and files
                        entry = entry.rstrip('/') + '/'
                        directories.append(entry)
                    else:
                        files.append(entry)

            # Sort directories and files separately
            directories.sort(key=lambda s: s.lower())
            files.sort(key=lambda s: s.lower())

            # Combine sorted directories and files
            sorted_entries = directories + files
            return sorted_entries

        except Exception as e:
            logger.error(f'Error listing files: {e}', exc_info=True)
            return []


    logger.debug(f'Starting action execution API on port {args.port}')
    run(app, host='0.0.0.0', port=args.port)
